
import numpy as np 
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec#分割子图
from pylab import mpl
import mplfinance.original_flavor as mpf
mpl.rcParams['font.sans-serif']=['SimHei']
mpl.rcParams['axes.unicode_minus']=False

class ploter():
    def plot_performance(self,df):
        font = {
            'family': 'SimHei',
            'weight': 'bold',
            'size': 12
        }
        plt.rc("font", **font)

        d1 = df[['累积收益率', '指数累积收益率']].copy()
        d1.plot(figsize=(15, 7))
        plt.title('均线交易回测与指数收益率对比图', size=15)
        plt.xlabel('')
        ax = plt.gca()
        ax.spines['right'].set_color('none')
        ax.spines['top'].set_color('none')
        plt.show()

class DefTypesPool():
    def __init__(self):
        self.routes={}
    #装饰器函数，将各种绘图函数注册到self.routes    
    def route_types(self,types_str):
        def decorator(f):
            self.routes[types_str]=f
            return f
        return decorator
    #根据注册路径输出已注册的绘图函数
    def route_output(self,path):
        function_val=self.routes.get(path)
        if function_val:
            return function_val
        else:
            raise ValueError(f"Route '{path}' has not been registered")


class MplTypesDraw():
    # 实例化DefTypesPool类，mpl为类变量
    mpl = DefTypesPool()

    # 画折线图
    @mpl.route_types("line")
    def line_plot(df_index, df_dat, graph):
        for key, val in df_dat.items():
            graph.plot(np.arange(0, len(val)), val, label=key, lw=1.0)

    # 画K线图
    @mpl.route_types("ochl")
    def ochl_plot(df_index, df_dat, graph):
        mpf.candlestick2_ochl(graph, df_dat['open'], df_dat['close'],
                              df_dat['high'], df_dat['low'], width=0.5,
                              colorup='r', colordown='g')

    # 画成交量图
    @mpl.route_types("bar")
    def nar_plot(df_index, df_dat, graph):
        graph.bar(np.arange(0, len(df_index)), df_dat['bar_red'], facecolor='red')
        graph.bar(np.arange(0, len(df_index)), df_dat['bar_green'], facecolor='green')


    @mpl.route_types(u"filltrade")
    def filltrade_plot(df_index, df_dat, graph):
        # 绘制filltrade图
        signal_shift = df_dat['signal'].shift(1)
        signal_shift.fillna(value=-1, inplace=True)  # 序列最前面的NaN值用-1填充
        list_signal = np.sign(df_dat['signal'] - signal_shift)
        bs_singal = list_signal[list_signal != 0]

        skip_days = False
        for kl_index, value in bs_singal.iteritems():  # iteritems以迭代器形式返回
            if (value == 1) and (skip_days == False):
                start = df_index.get_loc(kl_index)
                skip_days = True
            elif (value == -1) and (skip_days == True):
                end = df_index.get_loc(kl_index) + 1  # 加1用于匹配[start:end]选取到end值
                skip_days = False

                if df_dat['jdval'][end - 1] < df_dat['jdval'][start]:  # 赔钱显示绿色
                    graph.fill_between(np.arange(start, end), 0, df_dat['jdval'][start:end], color='green', alpha=0.38)
                    is_win = False
                else:  # 赚钱显示红色
                    graph.fill_between(np.arange(start, end), 0, df_dat['jdval'][start:end], color='red', alpha=0.38)
                    is_win = True
                graph.annotate('获利\n' if is_win else '亏损\n',
                               xy=(end, df_dat['jdval'].asof(kl_index)),
                               xytext=(df_dat['xytext'][0], df_dat['xytext'][1]),
                               xycoords='data',
                               va=df_dat['va'],  # 点在标注下方
                               textcoords='offset points',
                               fontsize=df_dat['fontsize'],
                               arrowprops=df_dat['arrow'])
        # 整个时间序列填充为底色blue 透明度alpha小于后标注区间颜色
        graph.fill_between(np.arange(0, len(df_index)), 0, df_dat['jdval'], color='blue', alpha=.08)
    
class MplVisualIf(MplTypesDraw):
    def __init__(self):
        MplTypesDraw.__init__(self)
        
    def fig_creat(self,**kwargs):
        if 'figsize' in kwargs.keys():
            self.fig=plt.figure(figsize=kwargs['figsize'],dpi=100,facecolor='white')
        else:
            self.fig=plt.figure(figsize=(14,7),dpi=100,facecolor='white')
        self.graph=self.fig.add_subplot(1,1,1)
        self.fig.autofmt_xdate(rotation=45)
    def fig_config(self,**kwargs):
        if 'legend' in kwargs.keys():
            self.graph.legend(loc=kwargs['legend'],shadow=True)
        if 'xlabel' in kwargs.keys():
            self.graph.set_xlabel(kwargs['xlabel'])
        else:
            self.graph.set_xlabel('日期')
        self.graph.set_title(kwargs['title'])
        self.graph.set_ylabel(kwargs['ylabel'])
        self.graph.set_xlim(0,len(self.index))
        
        if 'ylim' in kwargs.keys():
            bottom_lim=self.graph.get_ylim()[0]
            top_lim=self.graph.get_ylim()[1]
            range_lim=top_lim-bottom_lim
            self.graph.set_ylim(bottom_lim+range_lim*kwargs['ylim'][0],
                                top_lim+range_lim*kwargs['ylim'][1])
            
        if 'xticks' in kwargs.keys():
            self.graph.set_xticks(range(0,len(self.index),kwargs['xticks']))
        else:
            self.graph.set_xticks(range(0,len(self.index),15))
        if 'xticklabels' in kwargs.keys():
            self.graph.set_xticklabels([self.index.strftime(kwargs['xticklabels'])[index]\
                                        for index in self.graph.get_xticks()])
        else:
            self.graph.set_xticklabels([self.index.strftime('%Y%m%d')[index]\
                                        for index in self.graph.get_xticks()])
    def fig_show(self,**kwargs):
        plt.show()
        
    def fig_output(self,**kwargs):
        self.index=kwargs['index']
        self.fig_creat(**kwargs)
        for path,val in kwargs['draw_kind'].items():
            #print(f"输出{path}可视化图表：")
            view_function=self.mpl.route_output(path)
            view_function(self.index,val,self.graph)
        self.fig_config(**kwargs)
        self.fig_show(**kwargs)

class MultiGraphIf(MplTypesDraw):
    app=DefTypesPool()
    @app.route_types("ochl")
    def ochl_graph(data,sub_graph,df_dat=None):
        type_dict={'open':data.open,
                   'close':data.close,
                   'high':data.high,
                   'low':data.low}
        
        view_function=MplTypesDraw.mpl.route_output("ochl")
        view_function(data.index,type_dict,sub_graph)
        
    @app.route_types("ma")
    def sma_graph(data,sub_graph,periods):
        for val in periods:
            type_dict={'MA'+str(val):data.close.rolling(val).mean()}
            view_function=MplTypesDraw.mpl.route_output("line")
            view_function(data.index,type_dict,sub_graph)
            
    @app.route_types("vol")
    def vol_graph(data,sub_graph,df_dat=None):
        type_dict={'bar_red':np.where(data.open<data.close,data.vol,0),
                   'bar_green':np.where(data.open>data.close,data.vol,0)}
        view_function=MplTypesDraw.mpl.route_output("bar")
        view_function(data.index,type_dict,sub_graph)
    
    @app.route_types("macd")
    def macd_graph(data,sub_graph,df_dat=None):
        macd_dif=data.close.ewm(span=12,adjust=False).mean()-data.close.ewm(span=26,adjust=False).mean()
        macd_dea=macd_dif.ewm(span=9,adjust=False).mean()
        macd_bar=2*(macd_dif-macd_dea)
        
        type_dict={'bar_red':np.where(macd_bar>0,macd_bar,0),
                   'bar_green':np.where(macd_bar<0,macd_bar,0)}
        view_function=MplTypesDraw.mpl.route_output("bar")
        view_function(data.index,type_dict,sub_graph)
        
        type_dict={'macd_dif':macd_dif,
                   'macd_dea':macd_dea}
        view_function=MplTypesDraw.mpl.route_output("line")
        view_function(data.index,type_dict,sub_graph)
    
    @app.route_types("kdj")
    def kdj_graph(data,sub_graph,df_dat=None):
        low_list=data.low.rolling(9,min_periods=1).min()
        high_list=data.high.rolling(9,min_periods=1).max()
        rsv=(data.close-low_list)/(high_list-low_list)*100
        data['K']=rsv.ewm(com=2,adjust=False).mean()
        data['D']=data['K'].ewm(com=2,adjust=False).mean()
        data['J']=3*data['K']-2*data['D']
        
        type_dict={'K':data.K,
                  'D':data.D,
                  'J':data.J}
        view_function=MplTypesDraw.mpl.route_output("line")
        view_function(data.index,type_dict,sub_graph)
    
    def __init__(self, **kwargs):
        MplTypesDraw.__init__(self)
        self.fig = plt.figure(figsize=kwargs['figsize'], dpi=100, facecolor="white")#创建fig对象
        self.graph_dict = {}
        self.graph_curr = []

        try:
            gs = gridspec.GridSpec(kwargs['nrows'], kwargs['ncols'],
                                   left = kwargs['left'], bottom = kwargs['bottom'], right = kwargs['right'], top = kwargs['top'],
                                   wspace = kwargs['wspace'], hspace = kwargs['hspace'], height_ratios = kwargs['height_ratios'])
        except:
            raise Exception("para error")
        else:
            for i in range(0, kwargs['nrows'], 1):
                self.graph_dict[kwargs['subplots'][i]] = self.fig.add_subplot(gs[i, :])

    def graph_run(self, stock_data, **kwargs):
        # 绘制子图
        self.df_ohlc = stock_data
        for key in kwargs:
            self.graph_curr = self.graph_dict[kwargs[key]['graph_name']]
            for path, val in kwargs[key]['graph_type'].items():
                view_function = MultiGraphIf.app.route_output(path)
                view_function(self.df_ohlc, self.graph_curr, val)
            self.graph_attr(**kwargs[key])
        plt.show()

    def graph_attr(self, **kwargs):

        if 'title' in kwargs.keys():
            self.graph_curr.set_title(kwargs['title'])

        if 'legend' in kwargs.keys():
            self.graph_curr.legend(loc=kwargs['legend'], shadow=True)

        if 'xlabel' in kwargs.keys():
            self.graph_curr.set_xlabel(kwargs['xlabel'])

        self.graph_curr.set_ylabel(kwargs['ylabel'])
        self.graph_curr.set_xlim(0, len(self.df_ohlc.index))  # 设置一下x轴的范围
        self.graph_curr.set_xticks(range(0, len(self.df_ohlc.index), kwargs['xticks']))  # X轴刻度设定 每15天标一个日期

        if 'xticklabels' in kwargs.keys():
            self.graph_curr.set_xticklabels(
                [self.df_ohlc.index.strftime(kwargs['xticklabels'])[index] for index in
                 self.graph_curr.get_xticks()])  # 标签设置为日期

            # X-轴每个ticker标签都向右倾斜45度
            for label in self.graph_curr.xaxis.get_ticklabels():
                label.set_rotation(45)
                label.set_fontsize(10)  # 设置标签字体
        else:
            for label in self.graph_curr.xaxis.get_ticklabels():
                label.set_visible(False)

















    